import argparse
from configs.config import global_var


def load_predict_config(para):
    # 设置参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--CT_path', type=str, default=para + "CT_data", help='CT_slice_data location')
    # parser.add_argument('--CT_pt_path', type=str, default=para + "PointClouds/spine107_pts.npy",
    #                     help='CT_pt_data location')
    parser.add_argument('--model_load_path', type=str,
                        default=para + 'data/train_history/deepBasic/model/6D_model_loss0.000595634279306978.pth',
                        help='load trained model to tester')
    parser.add_argument('--refine_load_path', type=str,
                        default=para + 'data/train_history/SAC_均匀5+5/model/6D_model_loss0.005666051991283894.pth',
                        help='load trained model to tester')
    parser.add_argument('--DRR_test_path', type=str,
                        default=para + "data/CT/均匀采样/test/标准正位/DRR",
                        help='tester data location')
    parser.add_argument('--test_json_path', type=str,
                        default=para + "data/CT/均匀采样/test/标准正位/label.json",
                        help='tester data_label location')
    parser.add_argument('--mode', type=str,
                        default='reg', help='the mode of tester')
    parser.add_argument('--nums', type=int, default=1000, help='the quantity of tester data')
    parser.add_argument('--saving_path', type=str, default=global_var.saving_address + 'distribution',
                        help='saving result path')
    parser.add_argument('--saving_drr_path', type=str, default=global_var.saving_address + 'drr_vis_result',
                        help='saving result path')
    args = parser.parse_args()
    para_dict = {
        'CT_path': args.CT_path,
        # 'CT_pt_path': args.CT_pt_path,
        'model_load_path': args.model_load_path,
        'refine_load_path': args.refine_load_path,
        'DRR_test_path': args.DRR_test_path,
        'test_json_path': args.test_json_path,
        'mode': args.mode,
        'nums': args.nums,
        'saving_path': args.saving_path,
        'saving_drr_path': args.saving_drr_path
    }
    return para_dict
